package com.leetcode.partition9;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;

import com.leetcode.common.TreeNode;

public class LC863二叉树中所有距离为K的结点 {

    private int index = 0;

    public static void main(String[] args) {
        TreeNode root = new TreeNode(3, new TreeNode(5, new TreeNode(6), new TreeNode(2, new TreeNode(7), new TreeNode(4))),
            new TreeNode(1, new TreeNode(0), new TreeNode(8)));
        System.out.println(new LC863二叉树中所有距离为K的结点().distanceK(root, new TreeNode(5), 2));
    }
    
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        if (k == 0) return Arrays.asList(target.val);
        //将二叉树转换为图，用邻接表存储
        int n = 501, m = n - 1, start = target.val, level = 0;
        int[] head = new int[n], element = new int[2 * m], next = new int[2 * m];
        boolean[] visited = new boolean[n];
        for (int i = 0; i < n; i++) head[i] = -1;
        //建图
        preOrder(root, head, element, next);
        //以target为起点进行扩散
        LinkedList<Integer> queue = new LinkedList<>();
        List<Integer> answer = new ArrayList<>();
        queue.addLast(start);
        while (!queue.isEmpty()) {
            int size = queue.size();
            level++;
            for (int i = 0; i < size; i++) {
                int u = queue.pollFirst();
                visited[u] = true;
                for (int j = head[u]; j != -1; j = next[j]) {
                    int v = element[j];
                    if (!visited[v]) queue.addLast(v);
                    if (!visited[v] && level == k) answer.add(v);
                }
            }
            if (level == k) break;
        }
        return answer;
    }

    private void preOrder(TreeNode root, int[] head, int[] element, int[] next) {
        if (root == null) return;
        if (root.left != null) createEdge(root.val, root.left.val, head, element, next);
        if (root.right != null) createEdge(root.val, root.right.val, head, element, next);
        preOrder(root.left, head, element, next);
        preOrder(root.right, head, element, next);
    }
    
    private void createEdge(int u, int v, int[] head, int[] element, int[] next) {
        element[index] = v;
        next[index] = head[u];
        head[u] = index++;
        element[index] = u;
        next[index] = head[v];
        head[v] = index++;
    }
}
